package grpc_utils

import (
	"context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
	"strings"
)

var (
	userIdKey   = "userId"
	tenantIdKey = "tenantId"
)

func NewUserMetaDataContextUnaryServerInterceptor() grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {

		//接收client发送的handle消息
		md, ok := metadata.FromIncomingContext(ctx)
		newCtx := context.Background()

		if ok {

			//将userid的value取出
			if value, ok := md[strings.ToLower(userIdKey)]; ok {
				newCtx = context.WithValue(ctx, userIdKey, value[0])
			}

			//将tenantid的value取出
			if value, ok := md[strings.ToLower(tenantIdKey)]; ok {
				newCtx = context.WithValue(ctx, tenantIdKey, value[0])
			}

		}
		//我们想继续做时间统计，不让直接返回
		res, err := handler(newCtx, req)
		return res, err

	}
}

func NewUserMetaDataContextUnaryClientInterceptor() grpc.UnaryClientInterceptor {
	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {

		//对原有的 context 追加参数，metadata.Join
		send, _ := metadata.FromOutgoingContext(ctx)

		userId, ok := ctx.Value(userIdKey).(string)
		if !ok {
			userId = ""
		}
		tenantId, ok := ctx.Value(tenantIdKey).(string)
		if !ok {
			tenantId = ""
		}
		newMD := metadata.Pairs(userIdKey, userId, tenantIdKey, tenantId)

		ctx = metadata.NewOutgoingContext(ctx, metadata.Join(send, newMD))

		err := invoker(ctx, method, req, reply, cc, opts...)

		return err

	}
}
